import sys
from PyQt5.QtWidgets import QMainWindow, QApplication
from PyQt5.QtWidgets import QWidget
from PyQt5.QtWidgets import QVBoxLayout
from PyQt5.QtWidgets import QPushButton
from PyQt5.QtWidgets import QDesktopWidget
from PyQt5.QtWidgets import QComboBox
from PyQt5.QtWidgets import QHBoxLayout
from PyQt5.QtWidgets import QLineEdit
from PyQt5.QtWidgets import QDialog
from PyQt5.QtWidgets import QListView
from PyQt5.QtWidgets import QTabWidget
from PyQt5.QtWidgets import QAction
from PyQt5.QtWidgets import qApp
from manager.FridaManager import FridaManager
from manager.TaskManager import TaskManager
from view.LoggerView import LogDialog
from PyQt5.QtCore import QStringListModel
from PyQt5.QtWidgets import QAbstractItemView
import os
import importlib
from configobj import ConfigObj
import threading
from PyQt5.QtCore import QThread


class ProcessInfoDialog(QDialog):
    def __init__(self, frida_manager=None):
        super(ProcessInfoDialog, self).__init__()
        self.result = ""
        self.frida_manager = frida_manager
        self.result = ''
        v_box = QVBoxLayout()
        h_box = QHBoxLayout()
        self.list_view = QListView()
        h_box.addWidget(self.list_view)
        self.list_view.setEditTriggers(QAbstractItemView.NoEditTriggers)
        v_box.addLayout(h_box)

        h_box = QHBoxLayout()
        btn_refresh = QPushButton('刷新')
        btn_refresh.clicked.connect(self.refresh)
        btn_yes = QPushButton('确定')
        btn_yes.clicked.connect(self.yes)
        btn_no = QPushButton('取消')
        btn_no.clicked.connect(self.no)
        h_box.addWidget(btn_refresh)
        h_box.addWidget(btn_yes)
        h_box.addWidget(btn_no)
        v_box.addLayout(h_box)

        self.setLayout(v_box)
        self.setWindowTitle('进程信息')
        self.setFixedHeight(500)
        self.setFixedWidth(400)

        self.update()

    def yes(self):
        self.result = self.list_view.currentIndex().data()
        self.accept()

    def no(self):
        self.result = ''
        self.accept()

    def refresh(self):
        self.update()

    def update(self):
        processes = self.frida_manager.get_processes()
        list_model = QStringListModel()
        list_model.setStringList(processes)
        self.list_view.setModel(list_model)

    def get_result(self):
        return self.result


class MainWin(QMainWindow):

    def __init__(self, conf):
        super(MainWin, self).__init__()
        self.script_buffer = None
        self.device = None
        self.session = None
        self.platforms = []
        self.config = []
        self.servers = []
        self.tab_object = []
        self.script = []
        self.is_pause = True
        self.current_script_path = ''
        self.current_obj = None

        self.init_config(conf)
        self.log_out = LogDialog()

        self.frida_manager = FridaManager(self.log_out)
        self.frida_manager.init_device()

        v_box = QVBoxLayout()
        h_box = QHBoxLayout()

        self.cb_device_type = QComboBox()
        self.cb_device_type.addItems(self.platforms)
        self.cb_device_type.setCurrentText(self.platforms[0])
        self.cb_device_type.currentTextChanged.connect(self.device_change)
        h_box.addWidget(self.cb_device_type)

        self.cb_hook = QComboBox()
        self.cb_hook.addItem('附加')
        self.cb_hook.addItem('自启动')
        self.cb_hook.currentIndexChanged.connect(self.hook_module_change)
        h_box.addWidget(self.cb_hook)

        self.cb_process = QComboBox()
        self.cb_process.addItem('包名')
        self.cb_process.addItem('进程号')
        h_box.addWidget(self.cb_process)

        self.cb_connect = QComboBox()
        self.cb_connect.addItem('usb')
        self.cb_connect.addItem('wifi')

        self.le_process = QLineEdit()
        self.le_process.setText('com.hx.hxsocket')
        h_box.addWidget(self.le_process)

        self.btn_select = QPushButton('选择')
        self.btn_select.clicked.connect(self.select_click)
        self.btn_start_up = QPushButton('启动')
        self.btn_start_up.clicked.connect(self.start_up_click)
        self.btn_stop = QPushButton('停止')
        self.btn_stop.setEnabled(False)
        self.btn_stop.clicked.connect(self.stop_click)
        self.btn_pause = QPushButton('暂停')
        self.btn_pause.setEnabled(False)
        self.btn_pause.clicked.connect(self.pause_click)
        self.btn_start = QPushButton('开始')
        self.btn_start.setEnabled(False)
        self.btn_start.clicked.connect(self.start_click)
        h_box.addWidget(self.btn_select)
        h_box.addWidget(self.btn_start_up)
        h_box.addWidget(self.btn_stop)
        h_box.addWidget(self.btn_pause)
        h_box.addWidget(self.btn_start)
        v_box.addLayout(h_box)

        h_box = QHBoxLayout()
        self.tab = QTabWidget()
        self.tab.currentChanged.connect(self.tab_change)
        self.device_change()
        h_box.addWidget(self.tab)
        v_box.addLayout(h_box)

        main_frame = QWidget()
        main_frame.setLayout(v_box)

        self.add_menu_bar()
        self.add_status_bar()
        self.init_windows()
        self.setCentralWidget(main_frame)
        self.show()

    def select_click(self):
        dialog = ProcessInfoDialog(self.frida_manager)
        text = ''
        if dialog.exec_():
            if dialog.get_result():
                text = dialog.get_result()
        if text:
            if self.cb_process.currentText() == '进程号':
                items = text.split(' ')
                self.le_process.setText(items[0])
            else:
                items = text.split(' ')
                self.le_process.setText(items[1])

    def start_up_click(self):
        try:
            index = self.tab.currentIndex()
            items = self.tab_object[index]

            for key in items.keys():
                self.current_obj = items[key]
                device_type = self.cb_device_type.currentText()
                self.current_script_path = 'script/' + device_type + '/' + self.script[key]
                self.frida_manager.load_script(self.current_script_path)

            if self.cb_hook.currentText() == '附加':

                text = self.le_process.text()
                if self.cb_process.currentText() == '进程号':
                    text = int(text)
                self.frida_manager.attach_app(text)
                self.frida_manager.attach_load(self.current_obj.on_message)
            else:
                self.frida_manager.swap_app(self.le_process.text())
                self.frida_manager.swap_load(self.current_obj.on_message)

            self.btn_stop.setEnabled(True)
            self.btn_pause.setEnabled(True)
            self.btn_start_up.setEnabled(False)
        except Exception as e:
            self.log_out.exception_msg(self, str(e))

    def stop_click(self):
        try:
            self.frida_manager.detach()
            self.btn_stop.setEnabled(False)
            self.btn_pause.setEnabled(False)
            self.btn_start.setEnabled(False)
            self.btn_start_up.setEnabled(True)
        except Exception as e:
            self.log_out.exception_msg(self, str(e))

    def pause_click(self):
        try:
            self.btn_pause.setEnabled(False)
            self.btn_start.setEnabled(True)
            self.frida_manager.unload()
        except Exception as e:
            self.log_out.exception_msg(self, str(e))

    def start_click(self):
        try:
            self.btn_pause.setEnabled(True)
            self.btn_start.setEnabled(False)
            if self.cb_hook.currentText() == '附加':
                self.frida_manager.attach_load(self.current_obj.on_message)
            else:
                self.frida_manager.swap_load(self.current_obj.on_message)
        except Exception as e:
            self.log_out.exception_msg(self, str(e))

    def add_menu_bar(self):
        exit_action = QAction('&退出', self)
        exit_action.setShortcut('Ctrl+Q')
        exit_action.setStatusTip('退出程序')
        exit_action.triggered.connect(qApp.quit)

        log_action = QAction('&日志', self)
        log_action.setShortcut('Ctrl+R')
        log_action.setStatusTip('打开日志')
        log_action.triggered.connect(self.open_log_view)

        menu_bar = self.menuBar()
        tools = menu_bar.addMenu('&文件操作')
        tools.addAction(log_action)
        tools.addAction(exit_action)

        start_server64 = QAction('&启动frida64服务', self)
        start_server64.triggered.connect(self.start_server64)

        start_server32 = QAction('&启动frida32服务', self)
        start_server32.triggered.connect(self.start_server32)

        server = menu_bar.addMenu('&Frida服务')
        server.addAction(start_server64)
        server.addAction(start_server32)

    def add_tool_bar(self):
        exit_action = QAction('&退出', self)
        exit_action.setShortcut('Ctrl+Q')
        exit_action.triggered.connect(qApp.quit)
        toolbar = self.addToolBar('Exit')
        toolbar.addAction(exit_action)

    def add_status_bar(self):
        self.statusBar().showMessage('未连接')

    def hook_module_change(self):
        if self.cb_hook.currentText() == '自启动':
            self.cb_process.setCurrentText('包名')
            self.cb_process.setEnabled(False)
        else:
            self.cb_process.setEnabled(True)

    def device_change(self):
        self.tab.clear()
        text = self.cb_device_type.currentText()
        path = 'platforms.' + text + '.'
        items = self.config[text]
        for key in items.keys():
            try:
                module_path = path + key
                module = importlib.import_module(module_path)
                cls = getattr(module, key)
                obj = cls(self.frida_manager, self.log_out)
                self.tab_object.append({key: obj})
                self.tab.addTab(obj, items[key])
            except Exception as e:
                self.log_out.insert(str(e))

    def tab_change(self):
        self.btn_start_up.setEnabled(True)
        self.btn_stop.setEnabled(False)
        self.btn_pause.setEnabled(False)
        self.btn_start.setEnabled(False)
        print(self.tab.currentIndex())

    def init_windows(self):
        self.setGeometry(1000, 400, 1000, 700)
        self.setWindowTitle('Frida Tools')
        screen = QDesktopWidget().screenGeometry()
        form = self.geometry()
        x = (screen.width() - form.width()) / 2
        y = (screen.height() - form.height()) / 2
        self.move(int(x), int(y))

    def open_log_view(self):
        self.log_out.show()

    def init_config(self, conf):
        platforms = conf['platforms']
        for key in platforms.keys():
            self.platforms.append(platforms[key])
        self.config = conf
        self.servers = conf['servers']
        self.script = conf['script']

    def start_server64(self):
        try:
            path = self.servers['frida64']
            server = FridaServerThread(path)
            server.start()
        except Exception as e:
            self.log_out.insert(str(e))

    def start_server32(self):
        try:
            path = self.servers['frida32']
            server = FridaServerThread(path)
            server.start()
        except Exception as e:
            self.log_out.insert(str(e))

    def ios_proxy_forward(self):
        try:
            path = self.servers['frida32']
            server = FridaServerThread(path)
            server.start()
        except Exception as e:
            self.log_out.insert(str(e))


class FridaServerThread(threading.Thread):
    def __init__(self, path):
        super(FridaServerThread, self).__init__()
        self.path = path
        self.name = path[path.rfind('/') + 1:]

    def run(self):
        exc_cmd('adb forward tcp:27042 tcp:27042')
        exc_cmd('adb forward tcp:27043 tcp:27043')
        exc_cmd('adb shell su -c \' pkill -9 ' + self.name + '\'')
        exc_cmd('adb shell su -c \'' + self.path + '\'')


class IosProxyForward(threading.Thread):
    def __init__(self, path):
        super(IosProxyForward, self).__init__()

    def run(self):
        exc_cmd('iproxy 2222 22')


def exc_cmd(cmd):
    result = os.popen(cmd)
    res = result.read()
    result = []
    for line in res.splitlines():
        result.append(line)
    return result


def init():
    cur_path = os.path.dirname(os.path.realpath(__file__))
    cfg_path = os.path.join(cur_path, "config.ini")
    result = ConfigObj(cfg_path)
    return result


if __name__ == '__main__':

    try:
        config = init()
        app = QApplication(sys.argv)
        ex = MainWin(config)
        sys.exit(app.exec_())
    except Exception as e:
        print(e)
